#include "eigen/one_hot.h"

namespace  ldl_eigen
{

Eigen::MatrixXi OneHot::encode(const Eigen::VectorXi& labels, int num_classes) {
    if (labels.size() == 0) {
        return Eigen::MatrixXi(0, 0);
    }

    // 自动计算类别数（如果未指定）
    if (num_classes == -1) {
        num_classes = labels.maxCoeff() + 1;
    }

    // 检查有效性
    if (num_classes <= 0) {
        throw std::invalid_argument("num_classes must be positive");
    }

    if (labels.minCoeff() < 0 || labels.maxCoeff() >= num_classes) {
        throw std::out_of_range(
            "Labels must be in range [0, " + std::to_string(num_classes - 1) + "]");
    }

    // 初始化结果矩阵（全零）
    Eigen::MatrixXi one_hot(labels.rows(), num_classes);
    one_hot.setZero();

    // 设置对应位置为1
    for (int i = 0; i < labels.rows(); ++i) {
        one_hot(i, labels(i)) = 1;
    }

    return one_hot;
}

Eigen::VectorXi OneHot::decode(const Eigen::MatrixXi& one_hot) {
    if (one_hot.rows() == 0 || one_hot.cols() == 0) {
        return Eigen::VectorXi(0);
    }

    Eigen::VectorXi labels(one_hot.rows());

    for (int i = 0; i < one_hot.rows(); ++i) {
        int max_col = 0;
        one_hot.row(i).maxCoeff(&max_col);  // 找到值为1的列索引
        labels(i) = max_col;
    }

    return labels;
}

}